#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import pandas as pd
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'Arial', 'DejaVu Sans']
plt.rcParams['font.size'] = 16
plt.rcParams['axes.labelsize'] = 18
plt.rcParams['axes.titlesize'] = 20
plt.rcParams['legend.fontsize'] = 14
plt.rcParams['xtick.labelsize'] = 14
plt.rcParams['ytick.labelsize'] = 14
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.weight'] = 'bold'

df = pd.read_csv('timing_results.csv')

total_cols = [col for col in df.columns if col.startswith('total') and not col.startswith('total_all')]
main_cols = [col for col in df.columns if col not in ['method', 'test_size_MB'] and not col.startswith('total') and not col.startswith('clean')]

fig, axes = plt.subplots(2, 1, figsize=(18, 15))

def add_bar_labels_center(ax, data, cols):
    for i, bar in enumerate(ax.containers):
        for rect, method in zip(bar, df['method']):
            height = rect.get_height()
            if height > 10:
                col = cols[i]
                # 标签放在bar中间
                ax.text(
                    rect.get_x() + rect.get_width() / 2,
                    rect.get_y() + height / 2,
                    f"{col}\n{height:.1f}ms",
                    ha='center', va='center', fontsize=12, color='black', fontweight='bold'
                )

df.plot(x='method', y=total_cols, kind='bar', stacked=True, ax=axes[0])
axes[0].set_title('Total Stage Time (ms)')
axes[0].set_ylabel('Time (ms)')
axes[0].set_xlabel('Method')
axes[0].legend(
    loc='upper center',
    bbox_to_anchor=(0.5, 1.3),
    ncol=5,
    fontsize='small',
    frameon=True
)
add_bar_labels_center(axes[0], df, total_cols)
axes[0].set_xticklabels(df['method'], rotation=20, ha='right')

df.plot(x='method', y=main_cols, kind='bar', stacked=True, ax=axes[1])
axes[1].set_title('Detailed Stage Time (ms)')
axes[1].set_ylabel('Time (ms)')
axes[1].set_xlabel('Method')
axes[1].legend(
    loc='upper center',
    bbox_to_anchor=(0.5, 1.5),
    ncol=5,
    fontsize='small',
    frameon=True
)
add_bar_labels_center(axes[1], df, main_cols)
axes[1].set_xticklabels(df['method'], rotation=20, ha='right')

plt.tight_layout()
plt.savefig('timing_results.png', dpi=300, bbox_inches='tight')
print('Figure saved as timing_results.png')
plt.show()